# 改为使用 scipy 的 Lomb-Scargle 实现
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy.signal import lombscargle
from datetime import datetime

# 重新加载数据
df = pd.read_csv("/mnt/data/asteroid_photometry_results.csv")

# 筛选出 r 波段且 calmag 有效的观测点
df_r = df[(df['band'] == 'r') & (df['calmag'].notna())].copy()

# 转换时间为天数（相对时间）
df_r['timestamp'] = pd.to_datetime(df_r['obs_time'])
df_r['mjd'] = (df_r['timestamp'] - df_r['timestamp'].min()).dt.total_seconds() / 86400.0

# 准备数据
times = df_r['mjd'].values
magnitudes = df_r['calmag'].values
errors = df_r['magerr'].values

# 频率扫描范围
min_period = 0.1  # days
max_period = 5.0  # days
frequencies = np.linspace(1 / max_period, 1 / min_period, 10000)
angular_frequencies = 2 * np.pi * frequencies

# 使用 lombscargle 进行周期分析
magnitudes_mean_sub = magnitudes - np.mean(magnitudes)
power = lombscargle(times, magnitudes_mean_sub, angular_frequencies)

# 找到最优周期
best_frequency = frequencies[np.argmax(power)]
best_period = 1 / best_frequency

# 计算相位
phases = (times * best_frequency) % 1

# 绘图
import matplotlib.pyplot as plt
fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(8, 8))
ax1.plot(1 / frequencies, power)
ax1.set_xlabel("Period (days)")
ax1.set_ylabel("Power")
ax1.set_title(f"Detected Period: {best_period:.4f} days")

ax2.errorbar(phases, magnitudes, yerr=errors, fmt='o', markersize=5, alpha=0.7)
ax2.set_xlabel("Phase")
ax2.set_ylabel("Magnitude")
ax2.set_title("Folded Light Curve")
ax2.invert_yaxis()

plt.tight_layout()
plt.show()

best_period
